import re
import sys
import json
import configparser
import requests
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QPushButton, QLabel, QComboBox, QLineEdit, QMessageBox, \
    QInputDialog
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QFont
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

config = configparser.ConfigParser()
config.read('llm_config.ini', encoding='utf-8')
with open('model-url.json', 'r', encoding='utf-8') as file:
    model_details = json.load(file)

class CascadeComboBoxes(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setWindowTitle('Model Choose')
        self.setGeometry(300, 300, 300, 200)  # x, y, width, height

        self.layout = QVBoxLayout()

        font = QFont()
        font.setPointSize(10)

        self.comboBox1 = QComboBox()
        for i in range(len(list(model_details.keys()))):
            llm = list(model_details.keys())[i]
            self.comboBox1.addItems([llm])
        self.comboBox1.currentIndexChanged.connect(self.onFirstComboBoxChanged)
        self.comboBox1.setFont(font)

        self.comboBox2 = QComboBox()
        self.comboBox2.setFont(font)

        self.lineEdit = QLineEdit()
        self.lineEdit.setFont(font)

        self.button = QPushButton("run")
        self.button.clicked.connect(self.onButtonClicked)
        self.layout.addWidget(self.button)

        self.questionInput = QLineEdit()
        self.questionInput.setFont(font)
        self.questionInput.setPlaceholderText("输入您的问题")  # 设置占位符文本

        self.layout.addWidget(QLabel("Question:"), alignment=Qt.AlignRight)
        self.layout.addWidget(self.questionInput)

        self.layout.addWidget(QLabel("Model:"), alignment=Qt.AlignRight)
        self.layout.addWidget(self.comboBox1)
        self.layout.addWidget(QLabel("Version:"), alignment=Qt.AlignRight)
        self.layout.addWidget(self.comboBox2)
        self.layout.addWidget(QLabel("Temperatures:"), alignment=Qt.AlignRight)
        self.layout.addWidget(self.lineEdit)

        self.layout.setSpacing(10)
        self.layout.setContentsMargins(10, 10, 10, 10)

        self.applyStylesheet()

        self.setLayout(self.layout)

    def onFirstComboBoxChanged(self, index):
        self.comboBox2.clear()
        if index == 0:
            for i in range(len(list(model_details[list(model_details.keys())[index]].keys()))):
                model = list(model_details[list(model_details.keys())[index]].keys())[i]
                self.comboBox2.addItems([model])
        elif index == 1:
            for i in range(len(list(model_details[list(model_details.keys())[index]].keys()))):
                model = list(model_details[list(model_details.keys())[index]].keys())[i]
                self.comboBox2.addItems([model])
        elif index == 2:
            for i in range(len(list(model_details[list(model_details.keys())[index]].keys()))):
                model = list(model_details[list(model_details.keys())[index]].keys())[i]
                self.comboBox2.addItems([model])
        elif index == 3:
            for i in range(len(list(model_details[list(model_details.keys())[index]].keys()))):
                model = list(model_details[list(model_details.keys())[index]].keys())[i]
                self.comboBox2.addItems([model])
        elif index == 4:
            for i in range(len(list(model_details[list(model_details.keys())[index]].keys()))):
                model = list(model_details[list(model_details.keys())[index]].keys())[i]
                self.comboBox2.addItems([model])
        elif index == 5:
            for i in range(len(list(model_details[list(model_details.keys())[index]].keys()))):
                model = list(model_details[list(model_details.keys())[index]].keys())[i]
                self.comboBox2.addItems([model])

    def applyStylesheet(self):
        # 样式
        stylesheet = """
            QWidget {
                background-color: #F5F5F5;
                border: none;
            }
            QComboBox, QLineEdit {
                border: 1px solid #AAAAAA;
                border-radius: 5px;
                padding: 2px;
                min-width: 100px;
            }
            QComboBox QAbstractItemView {
                border: 1px solid #AAAAAA;
                border-radius: 5px;
                background-color: white;
            }
            QComboBox::drop-down {
                border: none;
                border-left: 1px solid #AAAAAA;
                border-radius: 0 5px 5px 0;
            }
            QComboBox:hover, QLineEdit:hover {
                border: 1px solid #888888;
            }
            QComboBox:focus, QLineEdit:focus {
                border: 1px solid #555555;
            }
            QPushButton {
                background-color: #4CAF50;
                border: none;
                border-radius: 5px;
                color: white;
                padding: 5px;
                min-width: 60px;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
            QPushButton:pressed {
                background-color: #3e8e41;
            }
            QLabel {
                color: #333333;
                font-weight: bold;
            }
            QWidget#CascadeComboBoxes {
                background-color: white;
                border: 1px solid #CCCCCC;
                border-radius: 5px;
            }
            """
        # 应用样式表
        self.setStyleSheet(stylesheet)

    def onButtonClicked(self):
        # 获取下拉菜单和输入框的值
        model = self.comboBox1.currentText()
        version = self.comboBox2.currentText()
        question = self.questionInput.text()
        temperature = self.lineEdit.text()
        while True:
            if type(temperature) == float and 0.1 <= temperature <= 1.0:
                break
            else:
                message_box = QMessageBox()
                message_box.setWindowTitle("Warning")
                message_box.setIcon(QMessageBox.Warning)
                message_box.setText("温度值必须在0.1到1.0之间")
                message_box.setInformativeText("重新输入")
                message_box.setStandardButtons(QMessageBox.Ok)
                if message_box.exec_() == QMessageBox.Ok:
                    text, ok = QInputDialog.getDouble(None, "温度：", "请输入温度值：", 0.1, 0.1, 1.0)
                    temperature = text
                    if type(temperature) == float and 0.1 <= temperature <= 1.0:
                        break
                message_box.exec_()
        if model in config:
            config.set('settings', 'current_model', model)
            config.set(model, 'model', version)
            config.set(model, 'temperature', str(temperature))
            config.set(model, 'base_url', model_details[model][version])
            with open('llm_config.ini', 'w') as configfile:
                config.write(configfile)
        result=llm_config('llm_config.ini', question)
        print(result)


def llm_config(path, question):
    config = configparser.ConfigParser()
    config.read(path, encoding='utf-8')
    current_model = config['settings']['current_model']
    model_config = config[current_model]
    print('chosen model name:', current_model, '\n', 'llm version:', model_config.get('model'), '\n', 'temperature:',
          model_config.get('temperature'), '\n', 'your question:', question)
    # if model_config.get('class_name') == 'ChatTongyi':
    #    model_choose = ChatOpenAI(
    #        base_url=model_config.get('base_url'),
    #        openai_api_key=model_config.get('openai_api_key')
    #        )
    # 星火
    if model_config.get('class_name') == 'ChatSparkLLM':
        model_choose = ChatSparkLLM(
            spark_api_url=model_config.get('base_url'),
            spark_app_id=model_config.get('APPID'),
            spark_api_key=model_config.get('APIKey'),
            spark_api_secret=model_config.get('APISecret'),
            spark_llm_domain=model_config.get('model'),
            streaming=False,
        )
        messages = [ChatMessage(
            role="user",
            content=question
        )]
        handler = ChunkPrintHandler()
        result = model_choose.generate([messages], callbacks=[handler])
        s = str(result)
        match = re.search(r"text=\'(.*?)\'", s)
        if match:
            extracted_text = match.group(1)
            #print(extracted_text)
            result=extracted_text

    elif model_config.get('class_name') == 'qianfan':
        params = {"grant_type": "client_credentials", "client_id": model_config.get("ak"),
                  "client_secret": model_config.get("sk")}
        response = requests.post("https://aip.baidubce.com/oauth/2.0/token", params=params)
        access_token = response.json().get("access_token")
        url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token={access_token}"
        headers = {'Content-Type': 'application/json'}
        payload = {
            "user_id": "python",
            "model": model_config.get('model'),
            "messages": [{"role": "user", "content": question}],
            "temperature": model_config.getfloat('temperature'),
            "disable_search": False,
            "enable_citation": False
        }
        json_payload = json.dumps(payload)

        response = requests.post(url, headers=headers, data=json_payload)
        result = response.json().get("result")
        #print(result)

    else:
        model_choose = ChatOpenAI(
            model=model_config.get('model'),
            temperature=model_config.getfloat('temperature'),
            base_url=model_config.get('base_url'),
            openai_api_key=model_config.get('openai_api_key')
        )
        chain = (
                model_choose
                | StrOutputParser()
        )
        result = chain.invoke(question)
        #print(result)
    return result

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = CascadeComboBoxes()
    ex.show()
    sys.exit(app.exec_())